30C - Shooting Gallery - CodeForces Solution

import sys
from types import GeneratorType

input = lambda: sys.stdin.buffer.readline().decode().rstrip()


def bootstrap(f):  
    stack = []
    def wrappedfunc(*args, **kwargs):
        if stack:
            return f(*args, **kwargs)
        else:
            to = f(*args, **kwargs)
            while True:
                if type(to) is GeneratorType:
                    stack.append(to)
                    to = next(to)
                else:
                    stack.pop()
                    if not stack:
                        break
                    to = stack[-1].send(to)
            return to
    return wrappedfunc

mod = 10**9 + 7
def f(i,j) : 
    if i == n :
        return 0 
    if dp[i][j] != -1:
        return dp[i][j]
    res = f(i+1,j) 
    if (arr[i][0] - arr[j][0])**2 + (arr[i][1] - arr[j][1])**2 <= (arr[i][2] - arr[j][2])**2 or j == -1 : 
                res = max(res, (arr[i][-1] + f(i+1,i) ))
    dp[i][j] = res
    return res

def solve():
    global n,arr,dp  
    dp = [[-1 for _ in range(1001)] for j in range(1001)]
    n = int(input())
    arr = []
    for _ in range(n) :
        arr.append( list(map(float,input().split())) )
    arr.sort( key = lambda x : x[2])
        print(f(0,-1))


 
    solve()